import torch
import glob
import cv2
from PIL import Image
from torchvision import transforms
import numpy as np
from resnet import ResNet


def load_image():
    IMAGE_PATH = "D:/vllm/cifar-10-batches-py"
    im_list = glob.glob("{}/test/*/*.png".format(IMAGE_PATH))
    np.random.shuffle(im_list)
    return im_list
    

label_name = ["airplane", "automobile", "bird", "cat", "deer",
              "dog", "frog", "horse", "ship", "truck"]

# 调整图片大小
test_transform = transforms.Compose([
    transforms.CenterCrop((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010)),
])

im_list = load_image()
net = ResNet()
net.load_state_dict(torch.load("cifar.model"))
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#net.to(device)
for im_path in im_list:
    net.eval()
    im_data = Image.open(im_path)

    inputs = test_transform(im_data)
    inputs = torch.unsqueeze(inputs, dim=0)

    #inputs = inputs.to(device)
    
    #分类图片，结果是 0..9
    outputs = net.forward(inputs)

    _, pred = torch.max(outputs.data, dim=1)
    print(label_name[pred.cpu().numpy()[0]])

    img = np.asarray(im_data)
    img = img[:, :, [1, 2, 0]]

    img = cv2.resize(img, (120, 120))
    cv2.imshow("img", img)
    cv2.waitKey(0)

